package com.lxm.kafka.test;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;

import java.util.List;
import java.util.Map;
import java.util.Random;

public class AuditPartitioner implements Partitioner {
    private Random random;

    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        String k = (String) key;
        List<PartitionInfo> partitionInfoList = cluster.availablePartitionsForTopic(topic);
        int partitionCount = partitionInfoList.size();
        int auditPartition = partitionCount - 1;
        int partition = random.nextInt(partitionCount);
        return key == null || k.isEmpty() || !k.contains("audit") ? partition : auditPartition;
    }

    /**
     * 实现必要资源的清理工作
     */
    @Override
    public void close() {

    }

    /**
     * 实现必要资源的初始话工作
     *
     * @param configs
     */
    @Override
    public void configure(Map<String, ?> configs) {
        random = new Random();
    }
}
